package com.xaicode.ctoroad.kafka;

import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.kafka.core.KafkaOperations;
import org.springframework.kafka.core.KafkaTemplate;
import org.springframework.stereotype.Component;
import org.springframework.util.concurrent.ListenableFuture;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.ExecutionException;

/**
 * @author Locker cjxia@isoftstone.com
 * @version 1.0
 */
@Slf4j
@Component
@SuppressWarnings("all")
public class KafkaSendHelper {

    private final KafkaTemplate<String, String> kafkaTemplate;

    private final AdminClient adminClient;

    private final KafkaConsumer<String, String> kafkaConsumer;

    /**
     * 是否开事务，
     * 开启事务时，只能使用带事务的操作 kafkaTemplate.executeInTransaction
     */
    @Value("${spring.kafka.enable-transaction}")
    private Boolean kafkaEnableTransaction;

    private int kafkaPartitions = 1;

    private int kafkaReplicationFactor = 1;

    public KafkaSendHelper(KafkaTemplate kafkaTemplate, AdminClient adminClient, KafkaConsumer<String, String> kafkaConsumer) {
        this.kafkaTemplate = kafkaTemplate;
        this.adminClient = adminClient;
        this.kafkaConsumer = kafkaConsumer;
    }

    /**
     * 发送单条消息
     *
     * @param topic   主题
     * @param message 消息体
     */
    public void send(String topic, String message) {
        if (kafkaEnableTransaction) {
            kafkaTemplate.executeInTransaction(
                    (KafkaOperations.OperationsCallback<String, String, ListenableFuture>) operations ->
                            kafkaTemplate.send(topic, message));
        } else {
            kafkaTemplate.send(topic, message);
        }
    }

    /**
     * 发送多条消息
     *
     * @param topic    主题
     * @param messages 多条消息，若要保持顺序需要使用 LinkedList
     */
    public void send(String topic, List<String> messages) {
        messages.forEach(message -> kafkaTemplate.send(topic, message));
    }

    /**
     * 创建主题，默认1个分区1个副本
     *
     * @param topic 主题
     */
    public void createTopic(String topic) {
        createTopic(topic, 1, (short) 1);
    }

    /**
     * 创建主题，指定分区和副本
     *
     * @param topic             主题
     * @param numPartitions     分区数量
     * @param replicationFactor 副本数量，不能多于broker个数
     */
    public void createTopic(String topic, int numPartitions, short replicationFactor) {
        NewTopic newTopic = new NewTopic(topic, numPartitions, replicationFactor);
        adminClient.createTopics(Collections.singletonList(newTopic));
    }

    /**
     * 创建新主题，跳过已存在主题，订阅新主题
     *
     * @param newTopics 主题集合
     */
    public void createAndSubscribeTopic(List<String> newTopics) {

        adminClient.listTopics().names()
                .whenComplete((currentTopics, throwable) -> {
                    log.info("kafkaConfig adminClient completed, current topics are : {}", currentTopics);

                    newTopics.removeIf(str -> !currentTopics.isEmpty() && currentTopics.contains(str));
                    if (!newTopics.isEmpty()) {
                        log.info("ready to create new kafka topics : {}", newTopics);
                        List<NewTopic> waitCreateTopicList = new ArrayList<>();
                        List<String> waitSubscribeTopicList = new ArrayList<>();
                        for (String ntp : newTopics) {
                            NewTopic newTopic = new NewTopic(ntp, kafkaPartitions, (short) kafkaReplicationFactor);
                            waitCreateTopicList.add(newTopic);
                            waitSubscribeTopicList.add(ntp);
                        }
                        adminClient.createTopics(waitCreateTopicList);
                    }
                });
    }


    public Set<String> listTopics() throws ExecutionException, InterruptedException {
        return adminClient.listTopics().names().get();
    }

    /**
     * TODO 订阅触发
     * @param topics
     * @return
     */
    public LinkedList<Map> subscribe(List<String> topics) {
        createAndSubscribeTopic(topics);
        kafkaConsumer.subscribe(topics);
        ConsumerRecords<String, String> records = kafkaConsumer.poll(Duration.ofMillis(100));

        LinkedList<Map> subContents = new LinkedList<>();
        for (int i = 0; i < 100; i++) {
            records.forEach(record -> {
                HashMap map = new HashMap();
                map.put("topic", record.topic());
                map.put("key", record.key());
                map.put("value", record.value());
                subContents.add(map);
                log.info("subscribe consumer message key-value : [{}] = [{}]", record.key(), record.value());
            });
        }

        return subContents;
    }

}
